use strict;
use warnings;
$| = 1;

use IO::File;
use SWISS::Entry;
use SWISS::ACs;
use SWISS::SQs;
use SWISS::DEs;
use SWISS::OSs;
use SWISS::GNs;
use SWISS::PE;
use SWISS::DTs;
use SWISS::IDs;

###############################################################################################
#						newUPver.pl  (version May 2015)
# Purpose: Download Uniprot distribution, divided by taxonomic category 
# 
# Uses program 'curl', which can be installed on Linux by 'sudo yum install curl'
# Uses SWISS::Entry; either install from CPAN, or untar provided scripts/SWISS_1.67.tar
#
# To use:
# 1. See "Change the following to fit requirements" below and change accordingly.
# 2. Go to the directory where you want the UniProt databases to be stored. 
#    You may use projects/DBfasta directory for this, but they can be elsewhere.
# 3. Copy newUPver.pl and SWISS_1_67.tar to this directory. Also copy newGOver.pl is you want GO annotations.
# 4. run 'tar xvf SWISS_1.67.tar'
# 5. Run 'perl newUPver.pl'
#
# Result:

# 1. Downloads Swissprot and Trembl .dat files for each taxonomic class and unpacks them to fasta files
# 2. If $do_full>0, creates the 'subset' DB, as described below
# 3. Creates TCW.anno for import into TCW
# Gzipped:
# The .fasta and .dat are not compressed. The .dat file needs to be uncompressed for newGOver.pl, but 
# can be removed after the GO MySQL TCW database is created by newGOver.pl. The .fasta files can be gzipped
# only if you plan to use diamond for searching.
# 

my $ftp = "ftp://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase";
my $ftpTax = "$ftp/taxonomic_divisions/";
my $ftpAll = "$ftp/complete/";
my $PAVE = "TCW.anno";  

############################################
## Note: you need plenty of memory and a fast download as UniProt has gotten really big
## Change the following to fit requirements

my $curDate = "Apr2015";  # Directory name for all UniProt taxonomic directories
                          # e.g. if you run this script from the projects/DBfasta directory,
                          # the UniProt subdirectories will be in projects/DBfasta/UniProt_Apr2015

# By downloading taxonomic databases separately, it allows TCW query by taxonomic database
# It makes blasting easier since its not against the whole UniProt at once.
# And you may not want the whole UniProt (see $do_all).
# Add (using the push command) or delete the taxononic databases you want
my @db;
push(@db, "plants");
push(@db, "fungi");
push(@db, "viruses");
push(@db, "invertebrates");
push(@db, "bacteria");
push(@db, "vertebrates");

my $trembl = 1; # set to 0 if you only want SwissProt downloaded

my $do_full = 1; # 0 - do not download the full SwissProt or Trembl
                 # 1 - download the full SwissProt but not Trembl. Make subset SwissProt
                 # 2 - download SwissProt and Trembl. Make subset SwissProt and Trembl
                 #     Note, if $trembl = 0, then that takes precedence
                 #     The entire Trembl is very big, so use 2 with caution. 
                 # Subset SwissProt is a fasta file of all SwissProt minus the entries 
                 # from the taxonomic databases. Similar for subset Trembl. 
                 # Hence, the full subset .fasta is dependent on what taxonomic databases are selected.
my $sp_full = "sp_full"; # directory for entire SwissProt
my $tr_full = "tr_full";   # directory for entire Trembl

##############################################
## MAIN
check_curl();
my $dir = "UniProt_$curDate";
mkchDir($dir);
  
my $suffix = &mkSuffix();
&make_LIB_cfg();
&make_Species_DBs();
&make_full_DBs();
exit(0);

#############################################################
sub make_LIB_cfg()
{
    my $ldir = `pwd`;
    chomp($ldir);
    print "Create $ldir/$PAVE\n";
    open OUT, ">$PAVE" or die $PAVE;
    my $db;
    my $j = 0;
    for (my $i = 0; $i< @db; $i++) 
    {
    	$db =$db[$i];
      $j++;
      print OUT "\nJPAVE_DBtaxo_$j = $db\n";
      print OUT "JPAVE_DBfasta_$j = $ldir/sp_$db/uniprot_sprot_" . $db . ".fasta\n";
      $j++;
      print OUT "\nJPAVE_DBtaxo_$j = $db\n";
      print OUT "JPAVE_DBfasta_$j = $ldir/tr_$db/uniprot_trembl_" . $db . ".fasta\n";
    }

    if ($do_full == 0) {return;}
    $db = "full_subset"
    $j++;
    print OUT "\nJPAVE_DBtaxo_$j = $db\n";
    print OUT "JPAVE_DBfasta_$j $ldir/sp_$db/uniprot_sprot" . $suffix . ".fasta\n";
   
    if ($trembl==0 || $do_full<2) {return;}
    $j++;
    print OUT "\nJPAVE_DBtaxo_$j = $db\n";
    print OUT "JPAVE_DBfasta_$j = $ldir/tr_$db/uniprot_trembl" . $suffix . ".fasta\n";
}

#################################################
sub make_Species_DBs 
{

  for (my $i = 0; $i< @db; $i++) 
  {
  	my $db =$db[$i];	
  	print "\n$db\n";

    # swissprot
  	mkchDir("sp_$db"); 
    my $fasta =  "uniprot_sprot_" . $db . ".fasta";
   	my $dat = "uniprot_sprot_" . $db . ".dat";
    my $gz  = $dat . ".gz";

    if (not -e $fasta) {
  	  	execute("curl -O " . $ftpTax . $gz);
  	  	execute("gunzip $gz");
  	  	sp2fasta($dat, $fasta);
    }
  	chdir("..");

	if ($trembl == 0) {return; }
	
	# trembl
	mkchDir("tr_$db");
	$fasta =  "uniprot_trembl_" . $db . ".fasta";
	$dat = "uniprot_trembl_" . $db . ".dat";
	$gz  = $dat . ".gz";

	if (not -e $fasta) {
	  	execute("curl -O " . $ftpTax . $gz);
		execute("gunzip $gz");
		sp2fasta($dat, $fasta);
	}
	chdir("..");
  }
}
###############################################
sub make_full_DBs 
{
  my ($fasta, $gz, $sfasta);
  
  if ($do_full == 0) {return;}
  
  # swissprot
  print "\n$sp_full\n";
  $fasta =  "uniprot_sprot.fasta";
  $gz  = $fasta . ".gz";
  mkchDir($sp_full);
  if (not -e $fasta) {
  	  execute("curl -O " . $ftpAll . $gz);
  	  execute("gunzip $gz");
  }
  if (@db > 0) {
      $sfasta = "uniprot_sprot_" . $suffix . ".fasta";
      createSubset("sp_", "_sprot_", $fasta, $sfasta);
  }
 
  chdir("..");

  #trembl
  if ($trembl==0 || $do_full<2) {return;}
  
  print "\n$tr_full\n";
  $fasta =  "uniprot_trembl.fasta";
  $gz  = $fasta . ".gz";
  mkchDir($tr_full);
  if (not -e $fasta) {
  	execute("curl -O " . $ftpAll . $gz);
  	execute("gunzip $gz");
  }
  if (@db > 0) {
      $sfasta = "uniprot_trembl_" . $suffix. ".fasta";
      createSubset("tr_", "_trembl_", $fasta, $sfasta);
  }
  chdir("..");
}

################################################################

sub sp2fasta
{
	my $inputfile = shift;
	my $outfile = shift;	
	print localtime()." Create fasta $outfile\n";

	my $fh = new IO::File $inputfile or 
	    die "Cannot open input file $inputfile: $!";

	open F, ">$outfile" or die $outfile;

	$/ = "\n\/\/";
	my $cnt=0;
	while(<$fh>) {
	    s/\r//g;
	    (my $entry_txt = $_) =~ s/^\s+//;
	    next unless $entry_txt;
	    $entry_txt .= "\n";
	    my $entry = SWISS::Entry->fromText( $entry_txt );
	    print F $entry->toFasta();
	    $cnt++;
	    if ($cnt % 10000) {print "   wrote " . $cnt . " fasta....\r";}
	}
	close F;
}
#########################################

sub createSubset
{
  	my $prefix = shift;
  	my $dbtype = shift;
	my $infile = shift;
	my $outfile = shift;
	my %ids;
	my $total=0;
	
 	print localtime()." Create subset $outfile\n";
	for (my $i = 0; $i< @db; $i++) {
		my $db = $db[$i];
		my $file = "../" . $prefix . $db . "/uniprot" . $dbtype . $db . ".fasta";	
		
		print " Reading $db $file\n";
		open IN, "<$file" or die "Cannot open $file";
		
		my $cnt=0;
		my $lines=0;
		# >sp|Q197F8|002R_IIV3 
		while(<IN>) {
			if (/^>\w\w\|\w*\|(\w*)/) {
				$ids{$1}++;
				$cnt++;
			}
			$lines++;
		}	
    	close IN;
    	print "   Read ", $cnt, " Lines ", $lines, "\n";
		$total += $cnt;
	}
 	print "Total to be removed: ", $total, "\n";
	if ($total<20) {
		print "***You will need to rerun this script to create the subset .fasta file.\n";
		print "   This script will not redo anything that has already been done\n";
		return;
	}
	print " Reading $infile\n";
	open IN, "<$infile" or die $infile;
	open OUT, ">$outfile" or die $outfile;
	my $prt = 0;
	my $cntRead=0;
	my $cntPrt=0;
	
	while(<IN>)
	{
		if (/^>\w\w\|\w*\|(\w*)/) 
		{	
			$cntRead++;
		 	if (exists $ids{$1}) {
				  $prt = 0;
			}
			else {
				  $prt = 1;
				  $cntPrt++;
				  print OUT;
			}
			if ($cntRead % 100000 == 0) {print "   Read $cntRead write $cntPrt...\r";}
		}	
		elsif ($prt) {print OUT;}
	}
	close IN;
	close OUT;
	print "   Read $cntRead write $cntPrt   \n";
}

##################################################################

sub check_curl
{
        my $curl = `which curl`;
        if (not $curl =~ /^\//)
        {
    	    print "Please install the curl utility for downloading Uniprot files\n";
    	    print "For example, 'sudo yum install curl'\n";
            exit(1);
        }
}

################################################################

sub execute
{
	my $cmd = shift;
    print localtime()." Exec $cmd\n";
    system($cmd); 
}

sub mkchDir
{
  my $dir = shift;
  if (-e $dir && -d $dir) {
    print "Directory $dir exists\n";
  }
  else {
    print "Create directory $dir\n";
    system("mkdir $dir");
  }
  chdir($dir) or die "Cannot chdir $dir";
}

sub mkSuffix {
  my $suffix ="_";
  for (my $i=0; $i < @db; $i++) {$suffix .= substr($db[$i],0,1);}
  if ($do_full) {print "All Uniprot subset suffix is: $suffix\n";}
  return $suffix;
}
